

# from langchain_community.utilities import SQLDatabase
#
# # 连接 MySQL 数据库
# db_user = "root"
# db_password = "root"
# db_host = "127.0.0.1"
# db_port = "3306"
# db_name = "spiders"
# db = SQLDatabase.from_uri(f"mysql+pymysql://{db_user}:{db_password}@{db_host}:{db_port}/{db_name}")
#
# print("哪种数据库：", db.dialect)
# print("获取数据表：", db.get_usable_table_names())
# # 执行查询
# res = db.run("SELECT count(*) FROM ali;")
# print("查询结果：", res)


from langchain_community.utilities import SQLDatabase
from langchain_openai import ChatOpenAI
from langchain.chains import create_sql_query_chain
from dotenv import load_dotenv
import os

load_dotenv()
# 连接 sqlite 数据库
# db = SQLDatabase.from_uri("sqlite:///demo.db")

# 连接 MySQL 数据库
db_user = "root"
db_password = "root"
db_host = "127.0.0.1"
db_port = "3306"
db_name = "spiders"
db = SQLDatabase.from_uri(f"mysql+pymysql://{db_user}:{db_password}@{db_host}:{db_port}/{db_name}")


# 加上大模型
# 创建模型对象
llm = ChatOpenAI(api_key=os.getenv("api_key"),
                 base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
                 model='qwen-plus'
                 )

chain = create_sql_query_chain(llm=llm, db=db)
# 限制使用的表
response = chain.invoke({"question": "有哪些城市招聘的岗位最多？", "table_names_to_use": ["ali"]})
print(response)
# 去除 SQLQuery
print(response[10:])
print("查询结果：", db.run(response[10:]))

